package com.mangosteen.interviewji.manager;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChoice;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Service
public class AiManager {

    @Resource
    private ArkService aiService;

    private final String DEFAULT_MODEL="deepseek-v3-250324";
    public String doChat(String userPrompt) {
        return doChat("",userPrompt,DEFAULT_MODEL);
    }

    public String doChat(String systemPrompt,String userPrompt) {
        return doChat(systemPrompt,userPrompt,DEFAULT_MODEL);
    }

    public String doChat(String systemPrompt,String userPrompt,String model) {
        // 初始化消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();

        // 创建用户消息
        ChatMessage systemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM) // 设置消息角色为用户
                .content(systemPrompt) // 设置消息内容
                .build();

        ChatMessage userMessage=ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userPrompt)
                .build();
        // 将用户消息添加到消息列表
        chatMessages.add(userMessage);
        chatMessages.add(systemMessage);

        // 创建聊天完成请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(model)// Replace with Model ID .
                .messages(chatMessages) // 设置消息列表
                .build();

        // 发送聊天完成请求并打印响应
        try {
            // 获取响应并打印每个选择的消息内容
            List<ChatCompletionChoice> choices=aiService.createChatCompletion(chatCompletionRequest)
                    .getChoices();
            if (choices!=null){
                return (String) choices.get(0).getMessage().getContent();
            }
        } catch (Exception e) {
            System.out.println("Error: " + e.getMessage());
        } finally {
            // 关闭服务执行器
            aiService.shutdownExecutor();
        }
        return "Error";
    }
}
